package kz.kstu.ds.signature;

import kz.kstu.ds.entity.SignedDocument;
import kz.kstu.ds.xml.JAXBHelper;
import sun.misc.BASE64Encoder;

import java.io.*;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;

/**
 * Date: 04.11.12
 * Time: 23:19
 *
 * @author Oleg Anastassov
 */
public class SignatureGenerator {

    // The private key used to sign input data.
    private PrivateKey privateKey;

    public SignatureGenerator() {
    }

    public SignatureGenerator(byte[] privateKeyBytes) throws GeneralSecurityException {
        initPrivateKey(privateKeyBytes);
    }

    private SignatureGenerator(String filePath) throws IOException, GeneralSecurityException {
        FileInputStream input = new FileInputStream(filePath);
        byte[] keyBytes = new byte[input.available()];
        input.read(keyBytes);
        input.close();
        initPrivateKey(keyBytes);
    }

    // Create the private key instance, using the key bytes passed in.
    private void initPrivateKey(byte[] privateKeyBytes) throws GeneralSecurityException {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("DSA");
        this.privateKey = keyFactory.generatePrivate(keySpec);
    }

    public byte[] signMessage(String message) throws NoSuchAlgorithmException, SignatureException, InvalidKeyException {
        Signature sig = Signature.getInstance("SHA1withDSA");
        sig.initSign(privateKey);
        sig.update(message.getBytes());
        return sig.sign();
    }

    public void sign(File in, File out, File key) {
        String content = null;
        try {
            BufferedReader br = new BufferedReader(new FileReader(in));
            StringBuffer sb = new StringBuffer();
            String line;
            while ((line = br.readLine()) != null) {
                sb.append(line);
            }
            content = sb.toString();
            br.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        try {
            SignatureGenerator sigGen = new SignatureGenerator(key.getName());
            byte[] message = sigGen.signMessage(content);
            BASE64Encoder encoder = new BASE64Encoder();
            String signature = encoder.encode(message);

            SignedDocument doc = new SignedDocument();
            doc.setContent(content);
            doc.setSignature(signature);

            if (!out.getName().endsWith(".xml")) {
                JAXBHelper.marshall(doc, new File(out.getName() + ".xml"));
            } else {
                JAXBHelper.marshall(doc, out);
            }

//            FileOutputStream output = new FileOutputStream(out);
//            output.write(signature.getBytes());
//            output.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
